﻿using TorchSharp;
using TorchSharp.Modules;
using static TorchSharp.torch;

namespace Qwen3.Module;

public class Qwen3RMSNorm : torch.nn.Module<Tensor, Tensor>
{
    private readonly float _eps;
    private readonly Parameter weight;

    public Qwen3RMSNorm(
        int hiddenSize,
        float eps = 1e-6f)
        : base(nameof(Qwen3RMSNorm))
    {
        this._eps = eps;
        this.weight = torch.nn.Parameter(torch.ones(hiddenSize));
        this.RegisterComponents();
    }

    public override Tensor forward(Tensor input)
    {
        using var _ = NewDisposeScope();
        var inputDtype = input.dtype;
        input = input.to_type(ScalarType.Float32);
        var variance = input.pow(2).mean([-1L], keepdim: true);
        var normed = input * torch.rsqrt(variance + this._eps);
        var output = this.weight * normed.to(inputDtype);
        return output.MoveToOuterDisposeScope();
    }
}